from langchain_openai import ChatOpenAI
from langchain.tools import BaseTool
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage,BaseMessage
from config import BackgroundAgentConfig
from pydantic import BaseModel, Field
from fastmcp import Client
from typing import Any, Dict
import asyncio


config = BackgroundAgentConfig()

class BackgroundAgent:
    def __init__(self,agent:str=None):
        self.agent = agent  # 代理实例
        self.agent_name = config.agent_name  # 代理名称
        # agent内容作为保留，留作之后调用其他代理

        self.api_key=config.api_key  # API密钥
        self.mcp_url=config.mcp_url  # MCP服务器URL
        self.available_tools = []  # 存储可用的MCP工具
        self.mcp_client = Client({
                "mcpServers": {
                    "background-mcp-server": {
                        "url": self.mcp_url,
                        "transport": "streamable-http"
                    }
                }
            })
        self.llm = None  # LangChain LLM实例
        self.llm_with_tools = None   # LLM与工具绑定
        self.agent_prompt = """
        你是一个专业的背景图片处理AI助手，专注于协助用户完成背景图片生成相关的任务。
        
        你的主要职责是：
        1. 准确理解用户提出的背景图片生成需求与描述
        2. 根据用户的描述，智能选择合适的工具生成高质量的背景图片
        3. 提供专业的建议和指导，确保生成的图片符合用户期望
        4. 对于复杂需求，可以组合使用多个工具。始终以友好、专业的态度回复用户。

        工作原则：
        - 始终以用户需求为中心，确保生成的图片满足用户期望
        - 在处理用户请求时，保持高效和准确
        - 当需求符合现有工具功能时，直接优先调用相应工具，避免不必要的复杂操作
        - 对于无法满足的需求，提供清晰的解释和建议
        - 始终保持专业和友好的沟通风格，确保用户体验良好
        - 对于工具的使用和对话的总结，始终要以现有的实际的内容进行回复，不可随意编造内容

        内容规范：
        - 对于与背景图片不相关或超出工具能力范围的请求，应当简洁明了的拒绝，并且指导用户如何调整请求
        - 回复的内容应当保持专业性，避免使用模糊或不准确的描述，避免口语化的表达
        - 不提供任何与背景图片生成无关的建议、信息和功能，专注于图片生成相关的任务
        """
        self.messages=[SystemMessage(content=self.agent_prompt)]  # 存储对话消息历史


    async def initialize(self):
        """异步初始化方法"""
        try:
            print("🔄 开始异步初始化...")
            self._initialize_llm()
            await self.find_tools()
            self.bind_tools_to_llm()
            print("✅ 初始化完成")
            return True
        except Exception as e:
            print(f"❌ 初始化失败: {e}")
            return False

    
    def _initialize_llm(self):
        """初始化 LangChain LLM"""
        try:
            print(f"🧠 初始化LLM: {config.llm_model_name}")
            self.llm = ChatOpenAI(
                api_key=config.api_key,
                base_url=config.llm_model_base_url,
                model=config.llm_model_name,
                temperature=0.7,  # 添加温度参数
                timeout=60,
                max_retries=3
            )
            print(f"✅ LLM初始化成功")
            print("🔍 测试 LLM 连接...")
            test_response = self.llm.invoke("你好")
            print(f"✅ LLM 连接测试成功: {test_response.content}")
            return True
            
        except Exception as e:
            print(f"❌ LLM初始化失败: {e}")
            return False
        
    async def find_tools(self):
        """发现可用的MCP工具并转换为LangChain工具"""
        try:
            print("🔍 发现MCP工具...")
            
            async with self.mcp_client:
                # 获取MCP工具列表
                tools_result = await self.mcp_client.list_tools()
                
                print(f"📋 发现 {len(tools_result)} 个MCP工具")
                
                for tool in tools_result:
                    print(f"🔧 处理工具: {tool.name}")
                    print(f"   描述: {tool.description}")
                    print(f"   输入模式: {tool.inputSchema}")
                    
                    # 转换为字典格式
                    langchain_tool = {
                        "type": "function",
                        "function": {
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool.inputSchema
                        }
                    }
                    
                    self.available_tools.append(langchain_tool)
                    print(f"✅ 成功转换工具: {tool.name}")
        except Exception as e:
            print(f"❌ 工具发现失败: {e}")

    def bind_tools_to_llm(self):
        """将发现的MCP工具绑定到LLM"""
        self.llm_with_tools= self.llm.bind_tools(self.available_tools)
    
    
    async def agent_call_tool(self, tool_name: str, params: Dict):
        """
        调用指定的MCP工具
        :param tool_name: 工具名称
        :param params: 工具参数
        :return: 工具调用结果
        """
        try:
            async with self.mcp_client as mcp_client:
                result = await mcp_client.call_tool(tool_name, params)
                # 处理不同类型的返回结果
                if isinstance(result, list) and len(result) > 0:
                    if hasattr(result[0], 'text'):
                        return result[0].text
                    else:
                        return str(result[0])
                else:
                    return str(result)
        except Exception as e:
            error_msg = f"调用工具失败: {str(e)}"
            print(f"❌ {error_msg}")
            return error_msg

    # 主要方法，解析用户消息并调用LLM和工具
    async def message_resolve(self, message: str):
        """
        处理用户消息，调用LLM和工具
        :param message: 用户输入消息
        :return: LLM响应或工具调用结果
        """
        # 构建用户消息
        human_message = HumanMessage(content=message)
        
        # 将消息添加到历史记录
        self.messages.append(human_message)
        
        try:
            print(f"📝 处理消息: {message}")
            
            # 调用LLM
            response = await self.llm_with_tools.ainvoke(self.messages)
            
            # 检查是否有工具调用
            if response.tool_calls:
                # 有工具调用，处理工具调用
                final_result = await self.tool_call_resolve(response)
                return final_result
            else:
                # 没有工具调用，直接返回LLM响应
                self.messages.append(response)
                return response.content
                
        except Exception as e:
            error_msg = f"处理消息失败: {str(e)}"
            print(f"❌ {error_msg}")
            return error_msg
        

    async def tool_call_resolve(self, llm_message: AIMessage) -> str:
        """
        解析工具调用请求，调用相应的MCP工具，加入到消息响应
        Args:
            llm_message (AIMessage): LLM的响应消息，包含工具调用请求
        """
        if not isinstance(llm_message, AIMessage):
            return "输入消息类型错误"
        
        tool_calls = llm_message.tool_calls
        if not tool_calls:
            return "没有工具调用请求"
        
        print(f"🔧 发现 {len(tool_calls)} 个工具调用")
        
        # 构建完整的消息链
        messages = self.messages.copy()
        messages.append(llm_message)
        
        # 依次处理每个工具调用
        for tool_call in tool_calls:
            # 修复：正确获取工具名称和参数
            if hasattr(tool_call, 'name'):
                # 如果是对象形式
                tool_name = tool_call.name
                params = tool_call.args
                tool_call_id = tool_call.id
            elif isinstance(tool_call, dict):
                # 如果是字典形式
                tool_name = tool_call.get('name') or tool_call.get('function', {}).get('name')
                params = tool_call.get('args') or tool_call.get('function', {}).get('arguments', {})
                tool_call_id = tool_call.get('id', 'call_1')
            else:
                print(f"⚠️ 未知的工具调用格式: {tool_call}")
                continue
            
            print(f"🔧 调用工具: {tool_name}，参数: {params}")

            # 调用MCP工具
            result = await self.agent_call_tool(tool_name, params)
            
            # 创建工具消息
            tool_message = ToolMessage(
                tool_call_id=tool_call_id,
                name=tool_name,
                content=str(result)
            )
            
            # 添加工具调用结果到消息链
            messages.append(tool_message)
            print(f"✅ 工具调用完成: {tool_name}")
        
        # 让LLM基于工具调用结果生成最终响应
        final_response = await self.llm_with_tools.ainvoke(messages)
        
        # 更新消息历史
        self.messages = messages
        self.messages.append(final_response)
        
        return final_response.content  # 返回字符串内容
    
    async def chat(self):
        """
        启动聊天循环，持续接收用户输入并处理
        """
        while True:
            user_input = input("用户: ")
            if user_input.lower() in ["exit", "quit"]:
                print("退出聊天")
                break
            
            response = await self.message_resolve(user_input)
            print(f"{self.agent_name}: {response}")



if __name__ == "__main__":
    agent = BackgroundAgent()
    asyncio.run(agent.initialize())  # 异步初始化
    asyncio.run(agent.chat())  # 启动聊天循环
